import pandas as pd
import numpy as np

class QLearningTable:
    def __init__(self, actions, learning_rate=0.1, reward_decay=0.9, e_greedy=0.9):
        self.actions = actions  # a list
        self.lr = learning_rate # 学习率
        self.gamma = reward_decay   # 奖励衰减
        self.epsilon = e_greedy     # 贪婪度
        self.q_table = pd.DataFrame(np.zeros((7 , len(actions))) , columns = actions)   # 初始 q_table

    # 选择动作，当前是Q-learning
    def choose_action(self, observation):

        if np.random.uniform() < self.epsilon:  # 选择 Q value 最高的 action

            state_action = self.q_table.iloc[observation ,  : ]

            action = np.random.choice(state_action[state_action == np.max(state_action)].index)
        else:
            action = np.random.choice(self.actions)
        return action


    def learn(self , s , a , r , s_):
        q_predict = self.q_table.loc[s, a]
        if s_ != 0 or s_ != 1:
            q_target = r + self.gamma * self.q_table.loc[s_, :].max()  # next state is not terminal
        else:
            q_target = r  # next state is terminal
        self.q_table.loc[s, a] += self.lr * (q_target - q_predict)  # update
        #print(self.q_table)
